module efield

  use dimensions


  complex, private   :: phi02DFT(jx,jz)  ! FT of potential of background Efield;
                                         ! initialized once for later use in "get_efield"

  complex, private   :: density1DFT(jz)  ! FT in z at the center (x=0), diagnostics
  complex, private   :: ExFT(jx,jz)      ! stored by "get_efield" for diagnostics only
  complex, private   :: EzFT(jx,jz)      ! stored by "get_efield" for diagnostics only
  complex, private   :: phi2DFT(jx,jz)   ! stored by "get_efield" for diagnostics only 


  complex, private   :: E0C(jz)          ! background E-field at the center
  complex, private   :: Phi0C(jz)        ! center potential of background E-field    
  complex, private   :: Phi0T            ! complex amplitude of background potential

  real,    private   :: E0z(jx, jz)      ! background E-field


contains

!--------------------------------------------------------------------------------------
subroutine efield_init(ExFilter, xmod)                ! inputs

! initialized private "phi02DFT", "Phi0T", "Phi0C", "E0C", and "E0z"
!
!--------------------------------------------------------------------------------------

  use fft_wrap, only   : fft_forward, fft_backward
  use parameters, only : phi0, kz, iE0Modulation, iunit

  implicit none

  real,    intent(in)  :: ExFilter(jx)
  complex, intent(in)  :: xMod(jx)

  integer              :: ix, iz
  integer              :: ikx

  complex              :: Phi1Dx(jx)
  complex              :: density1D(jz)
  complex              :: phi02D(jx,jz)



      !-- initialization of background field 

      phi0T = phi0
      do iz=1,jz
      E0C(iz) =  -iunit * kz * phi0 *exp ( iunit * kz * (iz-1) * dz  )
      end do
       
 
      call fft_forward(jz,E0C,E0C)
      E0C=E0C/float(jz)
      call dealias(jz,E0C)   ! in Fourier space

      Phi0C(:) = iunit * E0C(:) / kZvec(:)

      call fft_backward(jz,E0C,E0C)

      call fft_backward(jz,Phi0C,Phi0C)  ! back in x-space



      !-- modulate source --
      !
      !  modulate INITIAL source in x direction
      !  E0C(iz) =  -iunit * kz * phi0 *exp ( iunit * kz * (iz-1) * dz  )     
      !  so when x=SpeckleWidth/2, E0z is down by a factor of 0.00674
      !  E0C(iz) = -iunit * kz * phi0 *exp ( iunit * kz * z  )

      if (iE0Modulation==1) then
         do ix=1,jx
           E0z(ix,:)    = xMod(ix) * E0C(:) 
           phi02D(ix,:) = xMod(ix) * phi0 *exp ( iunit * kz * zvec(:)  ) 
        enddo
      else
         do ix=1,jx
           E0z(ix,:)    = E0C(:)
           phi02D(ix,:) = phi0 *exp ( iunit * kz * zvec(:)  )
         enddo
      endif  

      !-- take Fourier transform of initial potential, store result

      do iz=1,jz
         Phi1Dx(:) = phi02D(:,iz)       !    density2D(all x,kz)   complex Phi1Dx(jx)
         call fft_forward(jx,Phi1Dx,Phi1Dx)
         Phi1Dx=Phi1Dx/float(jx)
         call deAlias(jx,Phi1Dx)
         phi02DFT(:,iz) = Phi1Dx(:) * ExFilter(:) !   ph02DFT( all kx, all z)
      enddo

      do ikx=1,jx
         density1D(:) = phi02DFT(ikx,:)
         call fft_forward(jz,density1D,density1D)
         density1D=density1D/float(jz)
         call dealias(jz,density1D)
         phi02DFT(ikx,:) = density1D(:)   !   ph02DFT( all kx, all kz)
      enddo

end subroutine efield_init


!-------------------------------------------------------------------------------------
subroutine efield_compute_E(dfu, ExFilter, &                      ! inputs 
                            Ez2D, Ex2D, OmegaBounce)              ! outputs
!-------------------------------------------------------------------------------------



  use fft_wrap,   only : fft_forward, fft_backward
  use parameters, only : iunit

  implicit none



  !-- external variables --


  real,    intent(in)  ::  dfu(jz,jv,jx,jDirection) 
  real,    intent(in)  ::  ExFilter(jx)

  real,    intent(out) ::  Ez2D(jx,jz)
  real,    intent(out) ::  Ex2D(jx,jz)
  real,    intent(out) ::  OmegaBounce



  !-- internal variables --

  real     PhiCenter
  integer  ikx, ikz
  integer  ix, iz, iv, iDirection


  real     Phi1D(jx)
  complex  density1D(jz),  Phi1Dx(jx),  Phi1Dz(jz)
  complex  density2DFT(jx,jz)
  complex  density2D(jx,jz)

      
  !-- solve Poisson equation to find potential --

  do ix=1,jx

     do iz=1,jz
        density1D(iz)=0.0
        do iDirection=1,2
           do iv=1,jv
              density1D(iz) = density1D(iz) + dfu(iz,iv,ix,iDirection)*dv 
           end do
        enddo
    end do

    
    call fft_forward(jz, density1D, density1D)
    density1D = density1D/float(jz)
    call dealias(jz,density1D)

    density1D(1)=0.0                  ! subtract uniform background - for all x (?)
    density2D(ix,:) = density1D(:)    ! density2D(x,kz) 
                                      ! partial Fourier transform
                                      ! COMPLEX density2D(jx,jz), LOOKS OK
  enddo  !  end of ix loop

  do ikz=1,jz

     Phi1Dx(:) = density2D(:,ikz)       !    density2D(all x,kz)   complex Phi1Dx(jx)
     
     call fft_forward(jx, Phi1Dx, Phi1Dx)
     Phi1Dx=Phi1Dx/float(jx)
     call deAlias(jx, Phi1Dx)

     density2DFT(:,ikz) = Phi1Dx(:) * ExFilter(:)  !  density2DFT(all kx,kz), filtered in kx

     phi2DFT(:,ikz) = density2DFT(:,ikz) / (kXvec(:)**2 + kZvec(ikz)**2 )

  enddo


  phi2DFT = phi2DFT + phi02DFT          !!!!!!!!!    TOTAL POTENTIAL

  phi2DFT(1,:) = 0                      !  spatially uniform in x not allowed (03/21/2011)
        
  density1DFT(:) = density2D(jx/2,:)                !   diagnostic


!-- differentiate potential to find electric field -- 

  do ikx=1,jx                                       ! get Ez2D(x,z) complex ExFT(jx,jz)
     Phi1Dz(:) = -iunit * kZvec(:) * phi2DFT(ikx,:)
     call fft_backward(jz, Phi1Dz, Phi1Dz)
     EzFT(ikx,:) = Phi1Dz(:)                       ! Ez(kx, all z)
  end do

  do iz=1,jz
     Phi1Dx(:) = EzFT(:,iz)
     call fft_backward(jx, Phi1Dx, Phi1Dx)
     Ez2D(:,iz) = Phi1Dx(:)
  enddo	                                             !  got Ez(all x, all z)



  do ikx=1,jx   !  get  Ex2D(x,z)      
     Phi1Dz(:) = -iunit * kXvec(ikx) * phi2DFT(ikx,:)
     call fft_backward(jz, Phi1Dz, Phi1Dz)
     ExFT(ikx,:) = Phi1Dz(:)                         ! Ex(kx, all z)  complex ExFT(jx,jz)
  end do

  do iz=1,jz
     Phi1Dx(:) = ExFT(:,iz)
     call fft_backward(jx, Phi1Dx, Phi1Dx)
     Ex2D(:,iz) = Phi1Dx(:)                   
  enddo                                               ! got Ex(all x, all z)

  density1D(:) = Ez2D(jx/2 + 1,:)

  call fft_forward(jz, density1D, density1D)
  density1D = density1D/float(jz)
  call dealias(jz, density1D)

  ikz=2

  ! amplitude of TOTAL potential at speckle center of sin(kz*z + phase)
  PhiCenter = 2 * abs(density1D(ikz) / kZvec(ikz))

  OmegaBounce = kZvec(ikz) * sqrt(PhiCenter)
     
  return

end subroutine efield_compute_E
    

!-------------------------------------------------------------------------------------
subroutine efield_evolve_E0(xMod)                                   ! input
!-------------------------------------------------------------------------------------

  use parameters, only : phi0, kz, vPhase, dt, iE0Modulation, iWaveFrame, iunit

  implicit none
    
  complex, intent(in)      :: xMod(jx)

  complex                  :: phi02D(jx,jz)  ! generated, but never used
  real                     :: ExpKVdt
  integer                  :: ix


  !------------------------------------------


  ! in lab frame, E0z has angular frequency vPhase*kz
  if ( iWaveFrame==0 ) then
     ExpKVdt = exp( - iunit*dt*vPhase*kz  )
  else 
     ExpKVdt = 1
  end if
     
  phi0T    = ExpKVdt * phi0T
  phi0C(:) = ExpKVdt * phi0C(:)
  E0C(:)   = ExpKVdt * E0C(:)
         
  if (iE0Modulation==1) then               !  modulate  source in x direction
     do ix=1,jx   
        ! E0C(iz) =  -iunit * kz * phi0 *exp ( iunit * kz * (iz-1) * dz  )
        E0z(ix,:) =      xMod(ix) * E0C(:) 
        ! so when x=SpeckleWidth/2, E0z is down by a factor of 0.00674
        phi02D(ix,:) =   xMod(ix) * phi0 *exp ( iunit * kz * zvec(:)  ) 
        ! E0C(iz) = -iunit * kz * phi0 *exp ( iunit * kz * z  )
     enddo
  else
     do ix=1,jx
        E0z(ix,:) = E0C(:)
        phi02D(ix,:) =  phi0 *exp ( iunit * kz * zvec(:)  )
     enddo
  endif
    
    
end subroutine efield_evolve_E0

!-------------------------------------------------------------------------------------
subroutine efield_write_spectra()
!-------------------------------------------------------------------------------------

  use parameters, only : runname

  implicit none

  integer    :: ix, iz, ikz, ikx

  real       :: avEx, avEz, Eksq, phiksq

 !-- density1DFT(jz) is FT in z at the center (x=0)

  open(12,file = trim(runname)//".EfieldFTz.txt")

  write(12,"('% z-spectra E-field along center line')")
  write(12,"('% 1.kz   2.|E(kz)|^2   3.|Phi(kz)|^2')")
  write(12,*)

  do ikz=2,nkz/2  - nkz/6 -1
     Eksq   = abs(density1DFT(ikz)/kZvec(ikz)) **2 
     phiksq =  Eksq / kZvec(ikz)**2
     write(12,501) kZvec(ikz), Eksq, phiksq
  end do

  close(12)


  open(12, file = trim(runname)//".EfieldFTx.txt")

  write(12,"('% x-spectra of z-averaged E-field')")
  write(12,"('% 1.kx   2.<|Ex(kx)|^2>   3.<|Ez(kx)|^2>')")
  write(12,*) 

  do ikx=1,nkx/2  - nkx/6 -1
     avEx = 0
     avEz = 0
     do iz = 1, jz
        avEx = avEx + abs( ExFT(ikx,iz) )**2
        avEz = avEz + abs( EzFT(ikx,iz) )**2
     end do
     write(12,501) kxVec(ikx), avEx, avEz
  end do

  close(12)



  !open(26,file='PhiFT2_vs_kxkz.dat')   
  !write(26,*)'% kz\kx'
  !write(26,*)'NA', (  kxvec(ix), ix=2,jx/2 )
  !do iz=2,jz/2
  !   write(26,501) kzVec(iz), (    abs(phi2DFT(ix,iz))**2, ix=2,jx/2 )  !   phi2DFT(jx,jz)
  !enddo
  !close(26)
    

501  format(  6400(e12.6,2x)  )

end subroutine efield_write_spectra


!-------------------------------------------------------------------------------------
subroutine efield_write_vs_time(t)                ! inputs 
!-------------------------------------------------------------------------------------

  implicit none


  real, intent(in)        :: t

  real                    :: avEx2(jx)

  integer                 :: ix, iz, ikz, ikx
  logical, save           :: first_call = .true.


  if (first_call) then
     open(122,file='Ez_vs_kxtime.dat')
     write(122,*) '% t\kx'
     write(122,*) 'NA', ( kxVec(ikx), ikx = 1,nkx/2  - nkx/6 -1 )
     close(122)
     
     first_call = .false.
  end if

  do ikx=1,nkx/2  - nkx/6 -1
     avEx2(ikx) = 0
     do iz = 1, jz
        avEx2(ikx) = avEx2(ikx) + abs( EzFT(ikx,iz) )**2
     end do
  end do

  open(122,file='Ez_vs_kxtime.dat', position='append')
  write(122,*) t, ( avEx2(ikx), ikx = 1, nkx/2  - nkx/6 -1 )
  close(122)

end subroutine efield_write_vs_time

!-------------------------------------------------------------------------------------
subroutine efield_get_epsilon(epsilon)                ! output 
!-------------------------------------------------------------------------------------

  use parameters, only : nz, kz, iself

  implicit none

  complex, intent(out)    :: epsilon
  complex                 :: phiInt, phiTot


  phiInt = 2*density1DFT(1+nz)

  phiTot = phi0T*kz*kz
                 
  if (iself==1) then
     epsilon = phiTot / (phiInt + phiTot)
     !phi = 2*density1DFT(1+nz)/kz**2 + phi0T
  else
     epsilon = ( phiTot - phiInt ) / phiTot 
     ! = phiExternal/phiTotal = (phiTotal-phiInternal)/phiTotal
  end if


end subroutine efield_get_epsilon

!-------------------------------------------------------------------------------------
subroutine efield_get_E0z_center(E0zCenter)            ! output
!-------------------------------------------------------------------------------------

  real, intent(out)       :: E0zCenter(jz)

  E0zCenter = E0z(jx/2+1, :)

end subroutine efield_get_E0z_center

!-------------------------------------------------------------------------------------

  
end module
